from torch_geometric.nn import GATConv, SAGPooling, GCNConv
from torch_geometric.nn import global_mean_pool as gap, global_max_pool as gmp
from torch_geometric.nn import MessagePassing
from torch_geometric.nn.conv.gcn_conv import GCNConv
import torch.nn as nn
import torch
import torch.nn.functional as F
from typing import List, Tuple

from models.model_metrics import cal_metrics
from dataset_adapter import VarmisuseTask

# 本地库


class VarmisuseLinearOutputLayer(nn.Module):
    def __init__(
        self,
        out_features,
        task_type,
        max_node_per_graph=50,
        criterion=nn.CrossEntropyLoss(),
        metrics=cal_metrics,
        device="cpu",
    ):
        super(VarmisuseLinearOutputLayer, self).__init__()
        
        self.max_node_per_graph = max_node_per_graph
        
        self.varmisuse_base_linear=nn.Sequential(
            nn.Linear(out_features,out_features),
            nn.LayerNorm(out_features)
        )
        
        task_number=2 if task_type==VarmisuseTask.LocAndRepairBug else 1
        
        self.varmiuse_output=nn.Sequential(
            nn.Linear(out_features,task_number)
        )
        
        # self.varmisuse_layer = nn.Sequential(
            
        # )
        
        self.task_type=task_type
        self.criterion=criterion
        self.metrics = metrics
        self.device=device
        self.out_features=out_features

    def forward(self,
                output,
                slot_id,
                candidate_ids,
                candidate_masks,
                label=None,
                output_label=False,
                metric_functions=None,
                **kwargs):
        
        output=output.view(-1,self.max_node_per_graph,self.out_features)
        
        #print(output.size())
        output=self.varmisuse_base_linear(output)
        #print(output.size())
        result=self.varmiuse_output(output)
        #result=torch.sum(result,1)
        #print(output.size())
        bug_location_predict=None
        repair_predict=None
        
        if self.task_type==VarmisuseTask.LocBug:
            bug_location_predict=result[:,:,0]
            #print(bug_location_predict.size())
        elif self.task_type==VarmisuseTask.RepairBug:    
            repair_predict=result[:,:,0]
            #print(repair_predict.size())
        else:
            bug_location_predict=result[:,:,0]
            repair_predict=result[:,:,1]
            #print(bug_location_predict.size())
            #print(repair_predict.size())
        return bug_location_predict,repair_predict